# 专门负责旅游景点推荐的Agent
from python_a2a import A2AServer, run_server, TaskStatus, TaskState, AgentCard, AgentSkill
import requests
import json
import os

class TravelAgent(A2AServer):
    def __init__(self, agent_card):
        super().__init__(agent_card=agent_card)
        # 预定义的景点数据库，包含价格信息（人民币）
        self.attractions_db = {
            "beijing": {
                "outdoor": [
                    {"name": "天安门广场", "type": "历史景点", "weather_requirement": "晴天/多云", "price": 0, "free": True},
                    {"name": "颐和园", "type": "皇家园林", "weather_requirement": "晴天", "price": 30, "free": False},
                    {"name": "长城", "type": "古迹", "weather_requirement": "晴天", "price": 45, "free": False},
                    {"name": "北海公园", "type": "公园", "weather_requirement": "晴天/多云", "price": 10, "free": False},
                    {"name": "景山公园", "type": "公园", "weather_requirement": "晴天", "price": 2, "free": False}
                ],
                "indoor": [
                    {"name": "故宫博物院", "type": "博物馆", "weather_requirement": "任何天气", "price": 60, "free": False},
                    {"name": "国家博物馆", "type": "博物馆", "weather_requirement": "任何天气", "price": 0, "free": True},
                    {"name": "798艺术区", "type": "艺术区", "weather_requirement": "任何天气", "price": 0, "free": True},
                    {"name": "王府井", "type": "购物区", "weather_requirement": "任何天气", "price": 0, "free": True},
                    {"name": "三里屯", "type": "商业区", "weather_requirement": "任何天气", "price": 0, "free": True}
                ],
                "special_weather": [
                    {"name": "什刹海", "type": "湖区", "weather_requirement": "雪天", "special": "雪景很美", "price": 0, "free": True},
                    {"name": "香山", "type": "山景", "weather_requirement": "雪天", "special": "雪后红叶更美", "price": 10, "free": False}
                ]
            },
            "shanghai": {
                "outdoor": [
                    {"name": "外滩", "type": "城市景观", "weather_requirement": "晴天/多云", "price": 0, "free": True},
                    {"name": "豫园", "type": "古典园林", "weather_requirement": "晴天", "price": 40, "free": False},
                    {"name": "东方明珠", "type": "地标建筑", "weather_requirement": "晴天", "price": 180, "free": False}
                ],
                "indoor": [
                    {"name": "上海博物馆", "type": "博物馆", "weather_requirement": "任何天气", "price": 0, "free": True},
                    {"name": "新天地", "type": "商业区", "weather_requirement": "任何天气", "price": 0, "free": True}
                ]
            }
        }
        print(f"🏛️ TravelAgent 初始化完成")

    def handle_task(self, task):
        """处理景点推荐任务"""
        message_data = task.message or {}
        content = message_data.get("content", {})
        text = content.get("text", "")
        
        print(f"🏛️ TravelAgent收到任务: {text}")
        
        # 解析来自WeatherAgent的信息
        request_data = json.loads(text) if text.startswith('{') else {"query": text}
        
        # 提取预算信息 - 支持多种数据格式
        budget_info = request_data.get("budget_info", {})
        budget_constraints = request_data.get("budget_constraints", {})
        
        # 优先使用budget_constraints，然后回退到budget_info
        if budget_constraints:
            budget_amount = budget_constraints.get("total_budget", 0)
            # 根据金额判断状态
            if budget_amount >= 200:
                budget_status = "sufficient"
            elif budget_amount >= 20:
                budget_status = "limited"
            else:
                budget_status = "insufficient"
        else:
            budget_amount = budget_info.get("amount", 0)
            budget_status = budget_info.get("status", "unknown")
        
        print(f"🏛️ TravelAgent检测到预算: {budget_amount}元, 状态: {budget_status}")
        
        # 如果预算状态是"impossible"，直接拒绝推荐
        if budget_status == "impossible" or budget_amount < 20:  # 20元是最低可行预算（景山公园2元+交通+餐饮）
            response = {
                "agent": "TravelAgent",
                "budget_status": "insufficient",
                "budget_amount": budget_amount,
                "recommendations": [],
                "message": f"预算{budget_amount}元不足以支持任何旅游活动。建议预算至少100元以上才能有意义的旅游体验。",
                "total_attractions": 0
            }
        elif "weather_analysis" in request_data:
            # 这是来自WeatherAgent的分析结果
            city = request_data.get("city", "").lower()
            weather_analysis = request_data["weather_analysis"]
            suitable_for_outdoor = request_data.get("suitable_for_outdoor", False)
            temperature_info = request_data.get("temperature_info", {})
            
            recommendations = self._get_attractions_by_weather(city, weather_analysis, suitable_for_outdoor, temperature_info, budget_amount)
            
            response = {
                "agent": "TravelAgent",
                "city": city,
                "weather_considered": True,
                "budget_considered": True,
                "budget_amount": budget_amount,
                "recommendations": recommendations,
                "total_attractions": len(recommendations)
            }
        else:
            # 直接查询，没有天气信息
            city = request_data.get("city", "beijing").lower()
            recommendations = self._get_all_attractions(city, budget_amount)
            
            response = {
                "agent": "TravelAgent", 
                "city": city,
                "weather_considered": False,
                "budget_considered": True,
                "budget_amount": budget_amount,
                "recommendations": recommendations,
                "note": "建议先查询天气信息以获得更精准推荐"
            }

        task.artifacts = [{"parts": [{"type": "text", "text": json.dumps(response, ensure_ascii=False)}]}]
        task.status = TaskStatus(state=TaskState.COMPLETED)
        print(f"🏛️ TravelAgent回复: 推荐了{len(response['recommendations'])}个景点")
        
        # 发送数据流信息到UI
        try:
            import requests
            requests.post("http://localhost:8000/api/notify", json={
                "type": "data_flow",
                "step": {
                    "agent": "TravelAgent",
                    "description": f"景点推荐完成 - {city}",
                    "input": request_data,
                    "output": response,
                    "upstream": "CoordinatorAgent",
                    "downstream": "CoordinatorAgent"
                }
            }, timeout=2)
        except:
            pass
        
        return task

    def _get_attractions_by_weather(self, city, weather_analysis, suitable_for_outdoor, temperature_info, budget_amount=0):
        """根据天气条件和预算推荐景点"""
        recommendations = []
        city_data = self.attractions_db.get(city, self.attractions_db["beijing"])
        
        # 预算过滤函数
        def filter_by_budget(attractions, budget):
            if budget <= 0:
                return [attr for attr in attractions if attr.get("free", False)]
            return [attr for attr in attractions if attr.get("price", 0) <= budget]
        
        if suitable_for_outdoor and weather_analysis["outdoor_activities"] in ["适宜", "非常适宜"]:
            # 天气好，推荐户外景点，但要过滤预算
            outdoor_attractions = filter_by_budget(city_data["outdoor"], budget_amount)[:3]  # 取前3个可负担的
            for attraction in outdoor_attractions:
                attraction = attraction.copy()  # 避免修改原数据
                attraction["reason"] = f"天气{weather_analysis['outdoor_activities']}户外活动"
                attraction["priority"] = "高"
                attraction["budget_friendly"] = attraction.get("price", 0) <= budget_amount
                recommendations.append(attraction)
                
        # 总是包含一些室内选项，预算过滤
        indoor_attractions = filter_by_budget(city_data["indoor"], budget_amount)[:2]  # 取前2个可负担的
        for attraction in indoor_attractions:
            attraction = attraction.copy()  # 避免修改原数据
            attraction["reason"] = "备选室内活动"
            attraction["priority"] = "中" if suitable_for_outdoor else "高"
            attraction["budget_friendly"] = attraction.get("price", 0) <= budget_amount
            recommendations.append(attraction)
            
        # 特殊天气推荐，预算过滤
        if weather_analysis["outdoor_activities"] == "特殊适宜":
            special_attractions = filter_by_budget(city_data.get("special_weather", []), budget_amount)
            for attraction in special_attractions:
                attraction = attraction.copy()  # 避免修改原数据
                attraction["reason"] = "特殊天气景观"
                attraction["priority"] = "高"
                attraction["budget_friendly"] = attraction.get("price", 0) <= budget_amount
                recommendations.append(attraction)
                
        # 根据温度调整建议
        if temperature_info.get("temp"):
            temp = temperature_info["temp"]
            for rec in recommendations:
                if temp > 30:
                    rec["note"] = "建议早晚游览，避开高温时段"
                elif temp < 0:
                    rec["note"] = "注意保暖，建议缩短户外停留时间"
        
        # 如果没有任何推荐（预算太低），至少推荐免费景点
        if not recommendations and budget_amount <= 50:
            free_attractions = []
            for category in ["outdoor", "indoor", "special_weather"]:
                category_data = city_data.get(category, [])
                free_attractions.extend([attr for attr in category_data if attr.get("free", False)])
            
            if free_attractions:
                for attraction in free_attractions[:3]:  # 最多3个免费景点
                    attraction = attraction.copy()
                    attraction["reason"] = "免费景点，适合低预算游览"
                    attraction["priority"] = "低"
                    attraction["budget_friendly"] = True
                    recommendations.append(attraction)
        
        return recommendations

    def _get_all_attractions(self, city, budget_amount=0):
        """获取城市所有景点，按预算过滤"""
        city_data = self.attractions_db.get(city, self.attractions_db["beijing"])
        all_attractions = []
        
        # 预算过滤函数
        def filter_by_budget(attractions, budget):
            if budget <= 0:
                return [attr for attr in attractions if attr.get("free", False)]
            return [attr for attr in attractions if attr.get("price", 0) <= budget]
        
        # 过滤后添加景点
        outdoor_filtered = filter_by_budget(city_data["outdoor"], budget_amount)[:2]
        indoor_filtered = filter_by_budget(city_data["indoor"], budget_amount)[:2]
        
        for attraction in outdoor_filtered:
            attraction = attraction.copy()
            attraction["budget_friendly"] = attraction.get("price", 0) <= budget_amount
            all_attractions.append(attraction)
            
        for attraction in indoor_filtered:
            attraction = attraction.copy()
            attraction["budget_friendly"] = attraction.get("price", 0) <= budget_amount
            all_attractions.append(attraction)
        
        return all_attractions

if __name__ == "__main__":
    agent_card = AgentCard(
        name="Travel Expert Agent",
        description="专业旅游景点推荐Agent，基于天气条件提供个性化景点推荐",
        url="http://localhost:7004",
        version="1.0.0",
        skills=[
            AgentSkill(
                name="Attraction Recommendation",
                description="根据城市和天气条件推荐合适的景点"
            ),
            AgentSkill(
                name="Weather-based Planning",
                description="基于天气分析制定旅游计划"
            ),
            AgentSkill(
                name="Activity Categorization",
                description="将景点按照室内外、天气要求进行分类"
            )
        ]
    )

    travel_agent = TravelAgent(agent_card)
    
    print("🏛️ TravelAgent 启动于 http://localhost:7004")
    run_server(travel_agent, host="0.0.0.0", port=7004)
